import torch
import torch.nn as nn
import torch.nn.functional as F
from 封装ST_GCN.model.graph import Graph
from 封装ST_GCN.model.stgcn_block import STGCNBlock


class STGCN(nn.Module):
    def __init__(self, in_channels, num_class, graph_args, num_nodes,
                 edge_importance_weighting, use_residual=True, **kwargs):
        super().__init__()

        # 传入num_nodes参数到Graph类
        self.graph = Graph(**graph_args, num_nodes=num_nodes)
        if graph_args['strategy'] == 'adaptive':
            self.A = self.graph.A
        else:
            A = torch.tensor(self.graph.A, dtype=torch.float32, requires_grad=False)
            self.register_buffer('A', A)

        # build networks
        if graph_args['strategy'] == 'adaptive':
            spatial_kernel_size = self.A.size(0)
        else:
            spatial_kernel_size = self.A.size(0)
        temporal_kernel_size = 9
        kernel_size = (temporal_kernel_size, spatial_kernel_size)
        # 更新 self.data_bn 的输入维度
        self.data_bn = nn.BatchNorm1d(in_channels * num_nodes)
        kwargs0 = {k: v for k, v in kwargs.items() if k != 'dropout'}
        self.st_gcn_networks = nn.ModuleList((
            STGCNBlock(in_channels, 64, kernel_size, 1, residual=False, **kwargs0),
            STGCNBlock(64, 64, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(64, 64, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(64, 64, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(64, 128, kernel_size, 2, residual=use_residual, **kwargs),
            STGCNBlock(128, 128, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(128, 128, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(128, 256, kernel_size, 2, residual=use_residual, **kwargs),
            STGCNBlock(256, 256, kernel_size, 1, residual=use_residual, **kwargs),
            STGCNBlock(256, 256, kernel_size, 1, residual=use_residual, **kwargs)
        ))

        # initialize parameters for edge importance weighting
        if edge_importance_weighting:
            if graph_args['strategy'] == 'adaptive':
                self.edge_importance = nn.ParameterList([
                    nn.Parameter(torch.ones(self.A.size()))
                    for i in self.st_gcn_networks
                ])
            else:
                self.edge_importance = nn.ParameterList([
                    nn.Parameter(torch.ones(self.A.size()))
                    for i in self.st_gcn_networks
                ])
        else:
            self.edge_importance = [1] * len(self.st_gcn_networks)

        # fcn for prediction
        self.fcn = nn.Conv2d(256, num_class, kernel_size=1)

    def forward(self, x):
        # data normalization
        N, C, T, V, M = x.size()
        x = x.permute(0, 4, 3, 1, 2).contiguous()
        x = x.view(N * M, V * C, T)
        x = self.data_bn(x)
        x = x.view(N, M, V, C, T)
        x = x.permute(0, 1, 3, 4, 2).contiguous()
        x = x.view(N * M, C, T, V)

        # forward
        for gcn, importance in zip(self.st_gcn_networks, self.edge_importance):
            if isinstance(self.A, nn.Parameter):
                A = self.A
            else:
                A = self.A * importance
            x = gcn(x, A)

        # global pooling
        x = F.avg_pool2d(x, x.size()[2:])
        x = x.view(N, M, -1).mean(dim=1)

        # prediction
        x = self.fcn(x.unsqueeze(-1).unsqueeze(-1))
        x = x.squeeze(-1).squeeze(-1)

        return x